package com.yifeng.repo.controller.rocketmq.message;

import com.gomcarter.frameworks.base.common.AssertUtils;
import com.gomcarter.frameworks.base.exception.CustomException;
import com.google.common.base.Strings;
import com.talkyun.utils.revised.Looper;
import com.yifeng.repo.base.api.TimeRange;
import com.yifeng.repo.base.utils.common.BaseUtil;
import com.yifeng.repo.controller.rocketmq.configure.RocketmqProperties;
import com.yifeng.repo.controller.rocketmq.message.consume.SubscribePullConsumerLooper;
import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.acl.common.AclClientRPCHook;
import org.apache.rocketmq.acl.common.SessionCredentials;
import org.apache.rocketmq.client.consumer.AllocateMessageQueueStrategy;
import org.apache.rocketmq.client.consumer.DefaultLitePullConsumer;
import org.apache.rocketmq.client.consumer.rebalance.AllocateMessageQueueAveragely;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.common.consumer.ConsumeFromWhere;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.common.message.MessageQueue;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.function.Consumer;
import java.util.stream.Collectors;

/**
 * Created by daibing on 2023/8/5.
 */
@Slf4j
public class PullConsumeManager {
    protected final ConcurrentMap<String, Looper> topicSubExpression2PullConsumerLooper = new ConcurrentHashMap<>();
    protected final String applicationName;
    protected final RocketmqProperties properties;

    public PullConsumeManager(String applicationName, RocketmqProperties properties) {
        this.applicationName = applicationName;
        this.properties = properties;
    }

    public void init() {
        log.info("yfcloud controller PullConsumeManager init success: OWNER_ID={}, OWNER_NAME={}", BaseUtil.OWNER_ID, BaseUtil.OWNER_NAME);
    }

    public void destroy() {
        topicSubExpression2PullConsumerLooper.forEach((key, value) -> value.shutdown());
        log.info("yfcloud controller PullConsumeManager destroy success: OWNER_ID={}, OWNER_NAME={}", BaseUtil.OWNER_ID, BaseUtil.OWNER_NAME);
    }

    /**
     * 增加Topic消费（使用订阅模式：SubscribePullConsumerLooper）
     *
     * @param topic                        消息主题
     * @param subExpression                子表达式，比如：全部是 *，多个tag是 TagA || TagC || TagD
     * @param timeRangeList                指定时间区间
     * @param allocateMessageQueueStrategy 指定消费再平衡策略
     * @param msgListConsumer              拉取消息的消费者
     */
    public void addTopicConsumer(String topic, String subExpression, List<TimeRange> timeRangeList,
                                 AllocateMessageQueueStrategy allocateMessageQueueStrategy,
                                 Consumer<List<MessageExt>> msgListConsumer) {
        String key = getKey(topic, subExpression, allocateMessageQueueStrategy);
        Looper looper = topicSubExpression2PullConsumerLooper.get(key);
        if (looper != null) {
            return;
        }
        synchronized (topicSubExpression2PullConsumerLooper) {
            looper = topicSubExpression2PullConsumerLooper.get(key);
            if (looper != null) {
                return;
            }
            looper = buildSubscribePullConsumerLooper(topic, subExpression, timeRangeList, allocateMessageQueueStrategy, msgListConsumer);
            looper.start();
            topicSubExpression2PullConsumerLooper.put(key, looper);
        }
    }

    public void addTopicConsumer(String topic, String subExpression, List<TimeRange> timeRangeList,
                                 Consumer<List<MessageExt>> msgListConsumer) {
        this.addTopicConsumer(topic, subExpression, timeRangeList, new AllocateMessageQueueAveragely(), msgListConsumer);
    }

    public void addTopicConsumer(String topic, String subExpression,
                                 Consumer<List<MessageExt>> msgListConsumer) {
        List<TimeRange> timeRangeList = TimeRange.parse("00:00~23:59");
        this.addTopicConsumer(topic, subExpression, timeRangeList, new AllocateMessageQueueAveragely(), msgListConsumer);
    }

    /**
     * 移除Topic消费
     */
    public void removeTopicConsumer(String topic, String subExpression) {
        topicSubExpression2PullConsumerLooper.forEach((key, value) -> {
            if (key.startsWith(topic + "-" + subExpression)) {
                Looper looper = topicSubExpression2PullConsumerLooper.remove(key);
                if (looper != null) {
                    looper.shutdown();
                    log.info("removeTopicConsumer success: {}", looper.getName());
                }
            }
        });

    }

    /**
     * 移除Topic消费（指定再平衡策略）
     */
    public void removeTopicConsumer(String topic, String subExpression, AllocateMessageQueueStrategy allocateMessageQueueStrategy) {
        String key = getKey(topic, subExpression, allocateMessageQueueStrategy);
        Looper looper = topicSubExpression2PullConsumerLooper.remove(key);
        if (looper != null) {
            looper.shutdown();
            log.info("removeTopicConsumer success: {}", looper.getName());
        }
    }

    /**
     * 列出当前消息主题的消费subExpression（订阅模式）
     */
    public List<String> listConsumerSubExpression(String topic) {
        return topicSubExpression2PullConsumerLooper.keySet().stream()
                .map(s -> s.substring((topic + "-").length()))
                .collect(Collectors.toList());
    }

    /**
     * 查询指定topic的消息队列
     */
    public List<MessageQueue> listMessageQueue(String topic) {
        DefaultLitePullConsumer pullConsumer = getConsumerFromLooper(topic);
        try {
            return new ArrayList<>(pullConsumer.fetchMessageQueues(topic));
        } catch (MQClientException e) {
            throw new CustomException(e);
        }
    }

    /**
     * 获取队列的统计数据
     *
     * @param topic     消息主题
     * @param timestamp 开始时间，单位毫秒
     * @return queueId -> searchOffset
     */
    public Map<Integer, Long> listMessageQueueStats(String topic, long timestamp) {
        DefaultLitePullConsumer pullConsumer = getConsumerFromLooper(topic);
        Map<Integer, Long> mqStats = new HashMap<>();
        try {
            Collection<MessageQueue> messageQueues = pullConsumer.fetchMessageQueues(topic);
            for (MessageQueue messageQueue : messageQueues) {
                mqStats.put(messageQueue.getQueueId(), pullConsumer.offsetForTimestamp(messageQueue, timestamp));
            }
        } catch (MQClientException e) {
            throw new CustomException(e);
        }
        return mqStats;
    }

    protected DefaultLitePullConsumer getConsumerFromLooper(String topic) {
        AssertUtils.notEmpty(topicSubExpression2PullConsumerLooper.values(), "当前还没有初始化消费者，不能获取消息队列列表！");
        Looper looper = topicSubExpression2PullConsumerLooper.entrySet()
                .stream()
                .filter(entry -> entry.getKey().startsWith(topic + "-"))
                .findFirst()
                .orElseThrow(() -> new CustomException("获取不到Topic对应的消费者：" + topic)).getValue();
        return ((SubscribePullConsumerLooper) looper).getPullConsumer();
    }

    protected Looper buildSubscribePullConsumerLooper(String topic, String subExpression, List<TimeRange> timeRangeList,
                                                      AllocateMessageQueueStrategy allocateMessageQueueStrategy,
                                                      Consumer<List<MessageExt>> msgListConsumer) {
        String key = getKey(topic, subExpression, allocateMessageQueueStrategy);
        DefaultLitePullConsumer pullConsumer = buildSubscribePullConsumer(topic, subExpression, allocateMessageQueueStrategy);
        try {
            return new SubscribePullConsumerLooper(key + "-looper", 0, 1000,
                    topic, pullConsumer.fetchMessageQueues(topic),
                    pullConsumer, timeRangeList, msgListConsumer);
        } catch (MQClientException e) {
            throw new CustomException("buildSubscribePullConsumerLooper failed " + key, e);
        }
    }

    /**
     * 订阅方式：不需要考虑消息队列分配问题，支持自动重平衡
     * <p>
     * 通过Topic和Tag实现消息分类：不同消费者消费不同Tag，对于同一个Topic的不同Tag，如果RocketMQ消费者设置了相同的消费组，会导致消息消费混乱。
     * https://support.huaweicloud.com/bestpractice-hrm/hrm-bp-003.html
     */
    protected DefaultLitePullConsumer buildSubscribePullConsumer(String topic, String subExpression, AllocateMessageQueueStrategy allocateMessageQueueStrategy) {
        // 获取topic配置
        RocketmqProperties.TopicProperties topicProperties = getTopicProperties(properties, topic);
        String consumerGroup = getConsumerGroup(topicProperties, subExpression);

        // build pull consumer
        AclClientRPCHook rpcHook = new AclClientRPCHook(new SessionCredentials(topicProperties.getAccessKey(), topicProperties.getSecretKey()));
        DefaultLitePullConsumer pullConsumer = new DefaultLitePullConsumer(consumerGroup, rpcHook);
        pullConsumer.setInstanceName(applicationName + "-" + topic + "-pull-subscribe-" + subExpression + "-" + allocateMessageQueueStrategy.getName() + "-" + consumerGroup);
        pullConsumer.setNamesrvAddr(topicProperties.getNamesrvAddr());
        pullConsumer.setConsumeFromWhere(topicProperties.getConsumeFromWhere());
        pullConsumer.setPullBatchSize(topicProperties.getPullBatchSize());
        pullConsumer.setPullThreadNums(topicProperties.getSubscribePullThreadNums());
        pullConsumer.setAllocateMessageQueueStrategy(allocateMessageQueueStrategy);

        // consumer start.
        String key = getKey(topic, subExpression, allocateMessageQueueStrategy);
        try {
            pullConsumer.subscribe(topic, subExpression);
            pullConsumer.start();
            log.info("buildSubscribePullConsumer success: {}, consumerGroup={}, pullBatchSize={}, pullThreadNums={}",
                    key, pullConsumer.getConsumerGroup(), pullConsumer.getPullBatchSize(), pullConsumer.getPullThreadNums());
            return pullConsumer;
        } catch (MQClientException e) {
            throw new CustomException("buildSubscribePullConsumer failed " + key + "-" + consumerGroup, e);
        }
    }

    protected RocketmqProperties.TopicProperties getTopicProperties(RocketmqProperties properties, String topic) {
        return properties.getTopicPropertiesList()
                .stream()
                .filter(s -> s.getTopic().equals(topic))
                .findFirst()
                .orElseThrow(() -> new CustomException("没有Topic配置：" + topic));
    }

    /**
     * 指定订阅表达式使用指定消费组，如果没有指定订阅表达式，或者如果没有配置，都使用默认消费组
     */
    protected String getConsumerGroup(RocketmqProperties.TopicProperties topicProperties, String subExpression) {
        return Strings.isNullOrEmpty(subExpression) || topicProperties.getConsumerGroupMapping() == null
                ? topicProperties.getConsumerGroup()
                : topicProperties.getConsumerGroupMapping().getOrDefault(subExpression, topicProperties.getConsumerGroup());
    }

    private String getKey(String topic, String subExpression, AllocateMessageQueueStrategy allocateMessageQueueStrategy) {
        return topic + "-" + subExpression + "-" + allocateMessageQueueStrategy.getName();
    }

}
